package com.starhub.integration.sdk.baidu.core;

import lombok.Builder;
import okhttp3.MediaType;
import okhttp3.OkHttpClient;
import okhttp3.Request;
import okhttp3.RequestBody;
import okhttp3.Response;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.starhub.utils.Util;

import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
import java.util.stream.Stream;

/**
 * 百度向量化模型调度
 */
@Builder
public class BaiduClient {
    private OkHttpClient okHttpClient;
    private String apiHost;

    private String apiKey;
    private String secretKey;
    private String accessToken;

    public void setAccessToken(String accessToken) {
        this.accessToken = accessToken;
    }

    /**
     * 获取百度向量化模型的access token
     * 
     * @return access token
     * @throws IOException 如果请求失败
     */
    public String getAccessToken() throws IOException {
        String url = apiHost + "/oauth/2.0/token";

        Map<String, String> params = new HashMap<>();
        params.put("grant_type", "client_credentials");
        params.put("client_id", apiKey);
        params.put("client_secret", secretKey);

        String queryString = params.entrySet().stream()
                .map(entry -> entry.getKey() + "=" + entry.getValue())
                .reduce((a, b) -> a + "&" + b)
                .orElse("");

        Request request = new Request.Builder()
                .url(url + "?" + queryString)
                .get()
                .build();

        try (Response response = okHttpClient.newCall(request).execute()) {
            if (!response.isSuccessful()) {
                throw new IOException("获取access token失败，状态码: " + response.code());
            }

            String responseBody = response.body().string();
            JSONObject responseJson = JSON.parseObject(responseBody);

            if (Util.isNotEmpty(responseJson.get("error"))) {
                throw new IOException("获取access token失败: " + Util.null2String(responseJson.get("error_description"))
                        + Util.null2String(responseJson.get("error_description")));
            }

            return Util.null2String(responseJson.get("access_token"));
        } catch (IOException e) {
            throw new IOException("获取access token失败: " + e.getMessage());
        }
    }

    /**
     * 获取百度向量化模型的embedding
     * 
     * @param text 需要向量化的文本
     * @return 文本的向量表示
     * @throws IOException 当API调用失败或响应解析失败时抛出
     */
    public float[] getEmbedding(String text) throws IOException {
        if (Util.isEmpty(text)) {
            throw new IOException("输入文本不能为空");
        }

        String url = apiHost + "/rpc/2.0/ai_custom/v1/wenxinworkshop/embeddings/embedding-v1?access_token="
                + accessToken;

        Map<String, Object> requestParams = Map.of("input", List.of(text));

        RequestBody body = RequestBody.create(
                JSON.toJSONString(requestParams),
                MediaType.parse("application/json; charset=utf-8")
        );

        Request request = new Request.Builder()
                .url(url)
                .addHeader("Content-Type", "application/json")  
                .post(body)
                .build();

        try (Response response = okHttpClient.newCall(request).execute()) {
            if (!response.isSuccessful()) {
                throw new IOException(String.format("API请求失败，状态码: %d", response.code()));
            }

            String responseBody = response.body().string();
            if (Util.isEmpty(responseBody)) {
                throw new IOException("API返回响应为空");
            }

            return parseEmbeddingResponse(responseBody);
        } catch (IOException e) {
            throw new IOException("调用百度向量化API失败: " + e.getMessage(), e);
        }
    }

    /**
     * 解析百度向量化模型的embedding响应
     * 
     * @param responseBody API响应内容
     * @return 向量数组
     * @throws IOException 当响应解析失败时抛出
     */
    private float[] parseEmbeddingResponse(String responseBody) throws IOException {
        try {
            JSONObject jsonObject = JSON.parseObject(responseBody);

            // 检查错误信息
            String errorCode = jsonObject.getString("error_code");
            String errorMsg = jsonObject.getString("error_msg");
            if (Util.isNotEmpty(errorCode)) {
                throw new IOException(String.format("API返回错误: [%s] %s", errorCode, errorMsg));
            }

            // 获取并验证数据
            JSONArray dataArray = jsonObject.getJSONArray("data");
            if (dataArray == null || dataArray.isEmpty()) {
                throw new IOException("API返回数据为空");
            }

            JSONObject dataObject = dataArray.getJSONObject(0);
            JSONArray embeddingArray = dataObject.getJSONArray("embedding");
            if (embeddingArray == null || embeddingArray.isEmpty()) {
                throw new IOException("API返回的embedding数据为空");
            }

            // 转换为float数组
            float[] result = new float[embeddingArray.size()];
            for (int i = 0; i < embeddingArray.size(); i++) {
                result[i] = Util.getFloatValue(embeddingArray.get(i), 0.0f);
            }
            return result;
        } catch (Exception e) {
            throw new IOException("解析API响应失败: " + e.getMessage(), e);
        }
    }

    /**
     * 使用百度重排序模型对文本列表进行评分
     * 
     * @param texts 需要评分的文本列表
     * @param query 查询文本
     * @return 评分列表
     * @throws IOException 当API调用失败或响应解析失败时抛出
     */
    public List<Double> getRerankerScores(List<String> texts, String query) throws IOException {
        if (texts == null || texts.isEmpty() || query == null || query.trim().isEmpty()) {
            throw new IOException("输入文本或查询不能为空");
        }

        String url = apiHost + "/rpc/2.0/ai_custom/v1/wenxinworkshop/reranker/bce_reranker_base?access_token="
                + accessToken;

        Map<String, Object> requestParams = new HashMap<>();
        requestParams.put("query", query);
        requestParams.put("documents", texts);

        RequestBody body = RequestBody.create(
                JSON.toJSONString(requestParams),
                MediaType.parse("application/json; charset=utf-8")
        );

        Request request = new Request.Builder()
                .url(url)
                .addHeader("Content-Type", "application/json")
                .post(body)
                .build();

        try (Response response = okHttpClient.newCall(request).execute()) {
            if (!response.isSuccessful()) {
                throw new IOException(String.format("API请求失败，状态码: %d", response.code()));
            }

            String responseBody = response.body().string();
            if (Util.isEmpty(responseBody)) {
                throw new IOException("API返回响应为空");
            }

            return parseRerankerResponse(responseBody);
        } catch (IOException e) {
            throw new IOException("调用百度重排序API失败: " + e.getMessage(), e);
        }
    }

    /**
     * 解析百度重排序模型的响应
     * 
     * @param responseBody API响应内容
     * @return 评分列表
     * @throws IOException 当响应解析失败时抛出
     */
    private List<Double> parseRerankerResponse(String responseBody) throws IOException {
        try {
            JSONObject jsonObject = JSON.parseObject(responseBody);

            // 检查错误信息
            String errorCode = jsonObject.getString("error_code");
            String errorMsg = jsonObject.getString("error_msg");
            if (Util.isNotEmpty(errorCode)) {
                throw new IOException(String.format("API返回错误: [%s] %s", errorCode, errorMsg));
            }

            // 获取并验证数据
            JSONArray dataArray = jsonObject.getJSONArray("results");
            if (dataArray == null || dataArray.isEmpty()) {
                throw new IOException("API返回数据为空");
            }

            // 转换为Double列表
            List<Double> scores = new ArrayList<>();
            for (int i = 0; i < dataArray.size(); i++) {
                JSONObject result = dataArray.getJSONObject(i);
                scores.add(result.getDoubleValue("relevance_score"));
            }
            return scores;
        } catch (Exception e) {
            throw new IOException("解析API响应失败: " + e.getMessage(), e);
        }
    }
}